[#20567][perf] DFlash: snapshot the prefill context page counts at prepare time - #20568

Description

Fixes #20567.

With a DFlash drafter bound to the draft KV cache manager, _store_prefill_context reads the context rows' draft page counts with a blocking .tolist() of _ctx_block_counts at the end of every context forward, which waits for the forward's GPU work and exposes the executor's host work between prefill chunks as GPU idle.

With TLLM_DFLASH_PREFILL_CTX_SNAPSHOT=1 (new, read at import, default off):

• DFlashSpecMetadata.prepare() calls DFlashWorker._snapshot_ctx_block_counts(num_contexts): rows [0, num_contexts) of draft_kv_cache_block_offsets (already filled by attn_metadata.prepare() on this stream; nothing writes them before the refresh) give the page counts with the same formula as _refresh_ctx_block_tables; they are copied non-blocking into pinned memory and an event is recorded. The offsets tensor is the persistent buffer seen by the previous refresh (_ctx_offsets_src, new).
• _store_prefill_context takes the snapshot (_take_ctx_block_counts) and waits only on its event, which completes when the forward starts. It falls back to the device read when there is no snapshot or it does not match the forward: different row count, different offsets tensor, different draft KV manager or pool index, no block table. Never under CUDA graph capture; the snapshot is consumed once.

The values are identical to the device read (same formula, same tensor), so outputs do not change.

Performance (Rubin, TP8 across 2 nodes (MNNVL), Kimi K3 with a DFlash-style drafter, CUDA graphs, single request): the read went from up to about 190 ms of blocking to at most 0.12 ms; end to end, 128K TTFT dropped by about 0.36 s (the idle at chunk boundaries was about 11 ms per chunk before). Accuracy and needle checks unchanged.

#20550 also edits dflash.py, in different functions (decode-step prologue); the two apply independently.

Test Coverage

• New tests/unittest/_torch/speculative/test_dflash_prefill_ctx_snapshot.py (added to l0_b200.yml): over three successive steps (including empty rows and full rows), the snapshot equals the counts _refresh_ctx_block_tables computes from the same tensor, with the GPU still busy when it is taken; it is consumed once; it is rejected when the row count, offsets tensor, manager or pool changed; no snapshot under graph capture.
• tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py: the stub worker used by the two prepare() tests now binds _snapshot_ctx_block_counts and its state, and both tests run with the snapshot off and on.
• GPU results at 5fd9c02 (Rubin single GPU (200-SM engineering sample, SM107)):
  • test_dflash_prefill_ctx_snapshot.py: 6 passed (6 failed on main, as expected);
  • hw_agnostic -k dflash: the same pass / fail set as main (the existing failures there need the flash-attn package);
  • with TLLM_DFLASH_PREFILL_CTX_SNAPSHOT=1: two extra failures in test_dflash_dummy_slot.py, because its stub worker lacked the new method. Fixed in the test by the latest commit; the rerun is pending.
• Locally (no GPU): the snapshot / take logic runs with CPU stand-ins for pinned memory and events and gives the same counts and rejections. The same logic was validated earlier on Rubin against the device read (24 trials, 480 intervals, value-equal; four rejection cases; no snapshot with the flag off).
• pre-commit on the changed files: clean.

PR Checklist

• PR description clearly explains what and why.
• PR follows the TRT-LLM coding guidelines to the best of my knowledge.
• Test cases are provided for new code paths.
• No API changes, no new dependencies, no CODEOWNERS changes.
• Please check this after reviewing the above items as appropriate for this PR.